#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Bitcoin 区块链下载器
====================
这个程序允许下载完整的比特币区块链数据，支持断点续传、多线程下载和数据完整性验证。

作者: gA4ss
日期: 2025-06-21
"""

import os
import sys
import time
import json
import hashlib
import argparse
import logging
import requests
import threading
import urllib.parse
import shutil
import socket
import struct
from datetime import datetime
from concurrent.futures import ThreadPoolExecutor
from tqdm import tqdm

# 尝试导入BitTorrent相关库
try:
    import libtorrent as lt
    TORRENT_SUPPORT = True
except ImportError:
    TORRENT_SUPPORT = False
    logging.warning("未安装libtorrent库，BitTorrent下载功能将不可用。请使用pip install python-libtorrent安装")

# 尝试导入P2P网络相关库
try:
    import bitcoin
    from bitcoin.messages import msg_version, msg_verack, msg_getdata, msg_getblocks
    P2P_SUPPORT = True
except ImportError:
    P2P_SUPPORT = False
    logging.warning("未安装bitcoin-python库，P2P网络功能将不可用。请使用pip install bitcoin-python安装")

# 配置日志记录
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler("blockchain_downloader.log"),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger("blockchain_downloader")

# 比特币节点和镜像站点列表
BITCOIN_NODES = [
    "https://bitcoin.org/bin/",
    "https://bitcoincore.org/bin/",
    "https://bitcoincore.org/en/download/"
]

# 公共比特币API端点
BITCOIN_API_ENDPOINTS = [
    "https://blockchain.info",
    "https://api.blockcypher.com/v1/btc/main",
    "https://blockstream.info/api"
]

# 比特币Core端口
BITCOIN_CORE_PORT = 8333

# 比特币网络魔数
BITCOIN_NETWORK_MAGIC = bytes.fromhex('f9beb4d9')

# BitTorrent 追踪器列表
TORRENT_TRACKERS = [
    "udp://tracker.openbittorrent.com:80",
    "udp://tracker.publicbt.com:80",
    "udp://tracker.ccc.de:80",
    "udp://tracker.opentrackr.org:1337/announce"
]

# 比特币区块链数据种子文件
BLOCKCHAIN_TORRENT_URLS = {
    "full": "https://bitcointalk.org/index.php?topic=145386.0", # 完整区块链种子
    "pruned": "https://bitcointalk.org/index.php?topic=5262657.0" # 修剪区块链种子
}

class BlockchainDownloader:
    """比特币区块链下载管理器"""
    
    def __init__(self, output_dir, max_threads=4, use_torrent=False, verify=True, prune=False):
        """初始化下载管理器
        
        Args:
            output_dir (str): 区块链数据的输出目录路径
            max_threads (int): 最大并行下载线程数
            use_torrent (bool): 是否使用BitTorrent下载（更快但需额外依赖）
            verify (bool): 是否验证区块数据完整性
            prune (bool): 是否下载修剪版区块链
        """
        self.output_dir = os.path.abspath(output_dir)
        self.max_threads = max_threads
        self.use_torrent = use_torrent
        self.verify = verify
        self.prune = prune
        self.download_state = {}
        self.state_file = os.path.join(self.output_dir, ".download_state.json")
        self.torrent_handle = None
        self.torrent_session = None
        
        # 检查BitTorrent支持
        if self.use_torrent and not TORRENT_SUPPORT:
            logger.warning("由于缺少libtorrent库，切换到HTTP下载模式")
            self.use_torrent = False
        
        # 确保输出目录存在
        os.makedirs(self.output_dir, exist_ok=True)
        
        # 加载之前的下载状态（用于断点续传）
        self._load_state()
    
    def _load_state(self):
        """从状态文件加载之前的下载进度"""
        if os.path.exists(self.state_file):
            try:
                with open(self.state_file, 'r') as f:
                    self.download_state = json.load(f)
                logger.info(f"已加载下载状态: {len(self.download_state)} 个区块记录")
            except Exception as e:
                logger.error(f"加载下载状态失败: {e}")
                self.download_state = {}
    
    def _save_state(self):
        """保存当前的下载进度到状态文件"""
        try:
            with open(self.state_file, 'w') as f:
                json.dump(self.download_state, f, indent=2)
            logger.debug("下载状态已保存")
        except Exception as e:
            logger.error(f"保存下载状态失败: {e}")
    
    def start(self):
        """开始区块链下载过程"""
        logger.info("比特币区块链下载开始...")
        
        # 如果启用BitTorrent下载，使用BT下载方式
        if self.use_torrent and TORRENT_SUPPORT:
            logger.info("使用BitTorrent下载比特币区块链")
            return self.download_via_torrent()
        
        # 否则使用HTTP API下载
        logger.info("使用HTTP API下载比特币区块链")
        
        # 步骤1：获取区块链信息
        chain_info = self.get_blockchain_info()
        if not chain_info:
            logger.error("无法获取区块链信息，下载终止")
            return False
        
        # 步骤2：计算需要下载的区块范围
        start_height = 0
        end_height = chain_info['blocks']
        logger.info(f"区块链当前高度: {end_height}")
        
        # 步骤3：开始下载区块
        self.download_blocks(start_height, end_height)
        
        # 步骤4：验证下载的数据
        if self.verify:
            self.verify_blockchain()
        
        logger.info("区块链下载完成！")
        return True
        
    def download_via_torrent(self):
        """使用BitTorrent协议下载区块链"""
        try:
            # 初始化libtorrent会话
            self.torrent_session = lt.session()
            self.torrent_session.listen_on(6881, 6891)
            
            # 设置追踪器
            for tracker in TORRENT_TRACKERS:
                self.torrent_session.add_dht_router(*tracker.split("//")[1].split("/")[0].split(":"))
            
            # 选择种子类型
            torrent_type = "pruned" if self.prune else "full"
            torrent_url = BLOCKCHAIN_TORRENT_URLS[torrent_type]
            
            # 下载种子文件
            torrent_file = self._download_torrent_file(torrent_url)
            if not torrent_file:
                logger.error(f"无法获取种子文件从 {torrent_url}")
                return False
            
            # 设置参数
            params = {
                "save_path": self.output_dir,
                "storage_mode": lt.storage_mode_t.storage_mode_sparse,
                "auto_managed": True,
                "duplicate_is_error": True
            }
            
            # 添加种子
            self.torrent_handle = self.torrent_session.add_torrent(params)
            
            # 显示进度
            logger.info("开始下载区块链（BitTorrent模式）")
            print("注意：按 Ctrl+C 可暂停下载，您可以稍后使用 --resume 选项继续")
            
            # 下载监控循环
            pbar = tqdm(total=1000, desc="下载区块链", unit="%")
            prev_progress = 0
            
            while not self.torrent_handle.is_finished():
                status = self.torrent_handle.status()
                progress = int(status.progress * 1000)  # 准确到千分之一
                
                # 更新进度条
                if progress > prev_progress:
                    pbar.update(progress - prev_progress)
                    prev_progress = progress
                
                # 显示当前状态
                download_rate = status.download_rate / 1024 / 1024  # MB/s
                peers = status.num_peers
                state_str = str(status.state)
                if progress % 10 == 0:  # 减少输出频率
                    logger.info(f"下载进度: {status.progress*100:.2f}% | 速率: {download_rate:.2f} MB/s | 节点: {peers} | 状态: {state_str}")
                
                # 稍微暂停以减少CPU使用
                time.sleep(1)
            
            pbar.close()
            
            # 完成下载
            logger.info("种子下载完成！")
            
            # 验证数据
            if self.verify:
                self.verify_blockchain()
            
            return True
            
        except KeyboardInterrupt:
            logger.info("下载被用户中断，保存进度信息")
            # 这里保存种子状态信息以便继续
            return False
            
        except Exception as e:
            logger.error(f"BitTorrent下载出错: {str(e)}")
            logger.exception("错误详情:")
            return False
    
    def _download_torrent_file(self, url):
        """下载种子文件到本地"""
        try:
            # 如果文件已存在，直接返回
            torrent_file = os.path.join(self.output_dir, "blockchain.torrent")
            if os.path.exists(torrent_file):
                logger.info(f"使用现有种子文件: {torrent_file}")
                return torrent_file
            
            # 如果不存在，需要下载
            logger.info(f"下载种子文件从 {url}")
            
            # 注意：实际下载通常需要访问网页并解析正确的链接
            # 这里仅作演示，直接下载预先设置的URL
            response = requests.get("https://example.com/bitcoin-blockchain.torrent", timeout=30)
            
            if response.status_code == 200:
                with open(torrent_file, 'wb') as f:
                    f.write(response.content)
                logger.info(f"种子文件已下载到 {torrent_file}")
                return torrent_file
            else:
                logger.error(f"下载种子文件失败，状态码: {response.status_code}")
                return None
                
        except Exception as e:
            logger.error(f"下载种子文件时出错: {e}")
            return None
    
    def get_blockchain_info(self):
        """获取最新的区块链信息"""
        try:
            # 使用公共API获取区块链信息
            response = requests.get("https://blockchain.info/latestblock", timeout=10)
            if response.status_code == 200:
                data = response.json()
                return {
                    'blocks': data['height'],
                    'latest_hash': data['hash'],
                    'time': data['time']
                }
            else:
                logger.error(f"获取区块信息失败，状态码: {response.status_code}")
                return None
        except Exception as e:
            logger.error(f"请求区块链信息时出错: {e}")
            return None
    
    def download_blocks(self, start_height, end_height):
        """下载指定范围内的区块"""
        total_blocks = end_height - start_height
        logger.info(f"开始下载区块 {start_height} 到 {end_height} (共 {total_blocks} 个区块)")
        
        # 创建区块列表
        blocks_to_download = []
        for height in range(start_height, end_height + 1):
            # 检查是否已下载
            block_file = os.path.join(self.output_dir, f"blk{height//1000:05d}.dat")
            if os.path.exists(block_file) and str(height) in self.download_state:
                if self.download_state[str(height)]['status'] == 'complete':
                    continue
            
            blocks_to_download.append(height)
        
        if not blocks_to_download:
            logger.info("所有区块已下载完成")
            return True
        
        logger.info(f"需要下载 {len(blocks_to_download)} 个区块")
        
        # 使用线程池并行下载区块
        with ThreadPoolExecutor(max_workers=self.max_threads) as executor:
            # 使用tqdm显示进度
            list(tqdm(
                executor.map(self.download_single_block, blocks_to_download),
                total=len(blocks_to_download),
                desc="下载区块",
                unit="块"
            ))
        
        # 保存下载状态
        self._save_state()
        return True
    
    def download_single_block(self, height):
        """下载单个区块"""
        try:
            # 确定区块存储的文件
            file_index = height // 1000
            block_file = os.path.join(self.output_dir, f"blk{file_index:05d}.dat")
            
            # 确保目录存在
            os.makedirs(os.path.dirname(block_file), exist_ok=True)
            
            # 获取区块数据
            block_data = self.fetch_block_data(height)
            if not block_data:
                logger.error(f"无法获取区块 {height} 的数据")
                return False
            
            # 附加到文件
            with open(block_file, 'ab') as f:
                f.write(block_data)
            
            # 更新状态
            self.download_state[str(height)] = {
                'status': 'complete',
                'file': block_file,
                'time': time.time(),
                'size': len(block_data)
            }
            
            # 定期保存状态
            if height % 100 == 0:
                self._save_state()
                
            return True
        except Exception as e:
            logger.error(f"下载区块 {height} 时发生错误: {e}")
            self.download_state[str(height)] = {
                'status': 'error',
                'time': time.time(),
                'error': str(e)
            }
            return False
    
    def fetch_block_data(self, height):
        """从比特币网络获取特定高度的区块数据"""
        try:
            # 首先获取区块哈希
            hash_response = requests.get(f"https://blockchain.info/block-height/{height}?format=json", timeout=15)
            if hash_response.status_code != 200:
                logger.error(f"获取区块 {height} 哈希失败: {hash_response.status_code}")
                return None
            
            block_hash = hash_response.json()['blocks'][0]['hash']
            
            # 然后获取完整区块数据
            # 注意：实际应用中，应该直接从比特币P2P网络获取原始区块数据
            # 这里使用公共API仅用于演示目的
            block_response = requests.get(f"https://blockchain.info/rawblock/{block_hash}", timeout=15)
            if block_response.status_code != 200:
                logger.error(f"获取区块 {height} (hash:{block_hash}) 数据失败: {block_response.status_code}")
                return None
            
            # 在真实应用中，这里应返回原始二进制数据
            # 这里我们使用JSON响应来模拟，实际使用需将其转换为比特币的区块格式
            return block_response.content
        
        except Exception as e:
            logger.error(f"获取区块 {height} 数据时出错: {e}")
            return None
    
    def verify_blockchain(self):
        """验证下载的区块链数据完整性"""
        logger.info("开始验证区块链数据完整性...")
        
        # 1. 检查所有文件是否存在
        block_files = sorted([f for f in os.listdir(self.output_dir) if f.startswith("blk") and f.endswith(".dat")])
        logger.info(f"找到 {len(block_files)} 个区块文件")
        
        # 2. 简单验证文件大小和存在性
        for block_file in tqdm(block_files, desc="验证区块文件"):
            file_path = os.path.join(self.output_dir, block_file)
            if not os.path.exists(file_path):
                logger.error(f"缺少区块文件: {block_file}")
                continue
            
            file_size = os.path.getsize(file_path)
            if file_size == 0:
                logger.warning(f"区块文件 {block_file} 为空")
        
        # 3. 验证最后的区块哈希
        # 注意：完整验证需要解析区块数据并检查每个区块的哈希链接
        # 这超出了本程序的范围，需要更复杂的区块解析库
        
        logger.info("区块链数据验证完成")
        return True


def parse_args():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(
        description="比特币区块链下载工具 - 下载并验证完整的比特币区块链",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    
    parser.add_argument("-o", "--output-dir", 
                        default="./bitcoin_blockchain",
                        help="区块链数据的保存目录")
    
    parser.add_argument("-t", "--threads", 
                        type=int, 
                        default=4,
                        help="下载使用的线程数")
    
    parser.add_argument("--torrent", 
                        action="store_true",
                        help="使用BitTorrent下载（可能更快，但需要额外依赖）")
    
    parser.add_argument("--no-verify", 
                        action="store_true", 
                        help="跳过区块数据完整性验证")
    
    parser.add_argument("--resume", 
                        action="store_true",
                        help="从上次中断的地方继续下载")
    
    parser.add_argument("--blocks-only",
                        action="store_true",
                        help="只下载区块数据，不下载索引")
    
    parser.add_argument("--prune",
                        action="store_true",
                        help="下载修剪后的区块链（大小更小但功能受限）")
    
    return parser.parse_args()


def main():
    """主程序入口"""
    args = parse_args()
    
    # 检查必要的依赖
    if args.torrent and not TORRENT_SUPPORT:
        logger.warning("警告: 您选择了BitTorrent下载模式，但缺少libtorrent库")
        logger.warning("请使用以下命令安装必要的库： pip install python-libtorrent")
        logger.warning("切换到HTTP API下载模式...")
    
    # 显示计算机可用CPU核心信息
    cpu_cores = os.cpu_count() or 4
    if args.threads > cpu_cores:
        logger.warning(f"您设置的线程数 ({args.threads}) 超过CPU核心数 ({cpu_cores})")
        logger.warning(f"建议使用 {cpu_cores} 或更少的线程数")
    
    # 显示当前下载配置
    logger.info("=" * 50)
    logger.info("比特币区块链下载器")
    logger.info("=" * 50)
    logger.info(f"输出目录: {os.path.abspath(args.output_dir)}")
    logger.info(f"下载线程: {args.threads}")
    
    # 显示下载方式
    if args.torrent and TORRENT_SUPPORT:
        logger.info("下载方式: BitTorrent (更快，建议)")
    else:
        logger.info("下载方式: HTTP API (更慢，兼容)")
    
    logger.info(f"数据验证: {'\u5426' if args.no_verify else '\u662f'}")
    logger.info(f"断点续传: {'\u662f' if args.resume else '\u5426'}")
    logger.info(f"仅下载区块: {'\u662f' if args.blocks_only else '\u5426'}")
    logger.info(f"下载修剪版: {'\u662f' if args.prune else '\u5426'}")
    
    # 显示磅盘空间信息
    try:
        total, used, free = shutil.disk_usage(args.output_dir)
        logger.info(f"磅盘空间: 总计 {total // (1024**3)}GB, 已用 {used // (1024**3)}GB, 剩余 {free // (1024**3)}GB")
        
        # 检查磅盘空间是否足够
        required_space = 500 if not args.prune else 50  # 完整区块链约500GB，修剪版约50GB
        if free // (1024**3) < required_space:
            logger.warning(f"警告: 磅盘剩余空间不足! 需要 {required_space}GB, 但只有 {free // (1024**3)}GB")
            logger.warning("建议切换到较大的分区或启用--prune选项下载修剪后的区块链")
            user_continue = input("是否继续下载？(是/否) ")
            if user_continue.lower() not in ['y', 'yes', '是']:
                logger.info("用户取消下载")
                sys.exit(0)
    except Exception as e:
        logger.warning(f"无法获取磅盘空间信息: {str(e)}")
        
    logger.info("=" * 50)
    
    # 初始化下载器
    downloader = BlockchainDownloader(
        output_dir=args.output_dir,
        max_threads=args.threads,
        use_torrent=args.torrent and TORRENT_SUPPORT,
        verify=not args.no_verify,
        prune=args.prune
    )
    
    # 开始下载
    try:
        result = downloader.start()
        if result:
            logger.info("\n成功! 比特币区块链下载完成!")
            logger.info(f"数据保存在: {os.path.abspath(args.output_dir)}")
            if args.prune:
                logger.info("注意: 您下载了修剪版区块链，它只包含最近的区块头和部分交易数据")
        else:
            logger.error("下载失败，请检查日志文件了解详情")
            sys.exit(1)
            
    except KeyboardInterrupt:
        logger.info("\n\n下载已中断，保存当前进度...")
        downloader._save_state()
        logger.info("你可以使用以下命令继续下载:")
        logger.info(f"  python bitcoin_blockchain_downloader.py -o {args.output_dir} --resume {'--torrent' if args.torrent else ''} {'--prune' if args.prune else ''}")
        sys.exit(1)
        
    except Exception as e:
        logger.error(f"下载过程中出错: {e}")
        logger.exception("错误详情:")
        logger.info("你可以使用 --resume 选项稍后继续下载")
        sys.exit(2)


if __name__ == "__main__":
    main()
